import pandas as pd
import numpy as np
import mysql.connector
from mysql.connector import Error
from tqdm import tqdm
import os
from datetime import datetime


class DataProcessor:
    def __init__(self, db_config):
        self.data = pd.DataFrame()
        self.db_config = db_config
        self._load_data_from_db()
        self._preprocess_data()

    def _load_data_from_db(self):
        """从MySQL数据库加载数据"""
        try:
            connection = mysql.connector.connect(**self.db_config)
            if connection.is_connected():
                cursor = connection.cursor(dictionary=True)

                # 获取总记录数用于进度条
                cursor.execute("SELECT COUNT(*) as total FROM user_behavior")
                total_rows = cursor.fetchone()['total']
                print(f"开始从数据库加载数据，共 {total_rows} 条记录...")

                # 分块读取数据
                chunk_size = 100000
                offset = 0
                with tqdm(total=total_rows, unit='条') as pbar:
                    while offset < total_rows:
                        query = f"""
                        SELECT user_id, item_id, category_id, behavior_type, 
                               behavior_time as datetime 
                        FROM user_behavior 
                        LIMIT {chunk_size} OFFSET {offset}
                        """
                        cursor.execute(query)
                        chunk = cursor.fetchall()
                        if not chunk:
                            break

                        chunk_df = pd.DataFrame(chunk)
                        # 处理时间列
                        chunk_df['datetime'] = pd.to_datetime(chunk_df['datetime'])
                        self.data = pd.concat([self.data, chunk_df], ignore_index=True)

                        offset += chunk_size
                        pbar.update(min(chunk_size, total_rows - offset + chunk_size))

                print(f"数据加载完成，共加载 {len(self.data)} 条记录")

        except Error as e:
            print(f"数据库错误: {e}")
        finally:
            if 'connection' in locals() and connection.is_connected():
                cursor.close()
                connection.close()

    def _preprocess_chunk(self, chunk):
        """预处理数据块"""
        chunk['date'] = chunk['datetime'].dt.date
        chunk['hour'] = chunk['datetime'].dt.hour
        chunk['day_of_week'] = chunk['datetime'].dt.dayofweek
        chunk['is_buy'] = chunk['behavior_type'] == 'buy'
        return chunk

    def _preprocess_data(self):
        """预处理数据，筛选月度数据"""
        print("开始预处理数据：筛选11月数据...")
        self.data = self._preprocess_chunk(self.data)

        self.nov_data = self.data[
            (self.data['datetime'] >= '2017-11-01') &
            (self.data['datetime'] < '2017-12-01')
            ]
        print("11月数据筛选完成，开始筛选12月数据...")
        self.dec_data = self.data[
            (self.data['datetime'] >= '2017-12-01') &
            (self.data['datetime'] < '2018-01-01')
            ]
        print("12月数据筛选完成，开始提取购买数据...")
        self.nov_purchase = self.nov_data[self.nov_data['is_buy']]
        self.dec_purchase = self.dec_data[self.dec_data['is_buy']]
        print("数据预处理完成！")

    # 以下方法保持不变...
    def get_cumulative_purchases(self, months=None, year=2017):
        if months is None:
            months = [11, 12]

        result = {}
        for month in months:
            start_date = f'{year}-{month:02d}-01'
            if month == 12:
                end_date = f'{year + 1}-01-01'
            else:
                end_date = f'{year}-{month + 1:02d}-01'

            monthly_purchase = self.data[
                (self.data['datetime'] >= start_date) &
                (self.data['datetime'] < end_date) &
                (self.data['is_buy'])
                ]

            if monthly_purchase.empty:
                print(f'{year}-{month:02d} 没有购买数据')

            daily_purchase = monthly_purchase.groupby('date')['user_id'].count().reset_index(name='购买量')
            daily_purchase['累积购买量'] = daily_purchase['购买量'].cumsum()
            total = daily_purchase['购买量'].sum()

            result[f'{year}-{month:02d}_daily'] = daily_purchase
            result[f'{year}-{month:02d}_total'] = total

        return result

    def get_daily_visits(self):
        daily_visits = self.nov_data.groupby('date')['user_id'].count().reset_index(name='访问量')
        daily_visits['累积访问量'] = daily_visits['访问量'].cumsum()
        return daily_visits

    def get_category_purchase_ranking(self, year=2017, month=11):
        start_date = f'{year}-{month:02d}-01'
        if month == 12:
            end_date = f'{year + 1}-01-01'
        else:
            end_date = f'{year}-{month + 1:02d}-01'

        monthly_purchase = self.data[
            (self.data['datetime'] >= start_date) &
            (self.data['datetime'] < end_date) &
            (self.data['is_buy'])
            ]

        user_category_count = monthly_purchase.groupby('user_id')['category_id'].nunique().reset_index(
            name='购买商品品类数'
        )
        user_category_count = user_category_count.sort_values(
            by='购买商品品类数', ascending=False
        ).reset_index(drop=True)
        user_category_count['排名'] = user_category_count.index + 1
        return user_category_count

    def get_top_purchase_users(self, percentage=30):
        user_purchase_count = self.data[self.data['is_buy']].groupby('user_id')['item_id'].count().reset_index(
            name='购买数量')
        threshold = user_purchase_count['购买数量'].quantile(1 - percentage / 100)
        top_users = user_purchase_count[user_purchase_count['购买数量'] >= threshold].sort_values(
            by='购买数量', ascending=False
        ).reset_index(drop=True)
        return top_users

    def get_users_with_long_interval(self):
        user_purchase_times = self.data[self.data['is_buy']].sort_values(['user_id', 'datetime'])
        user_purchase_times['prev_time'] = user_purchase_times.groupby('user_id')['datetime'].shift(1)
        user_purchase_times['time_diff'] = (user_purchase_times['datetime'] - user_purchase_times['prev_time']).dt.days
        long_interval_users = user_purchase_times[user_purchase_times['time_diff'] > 5]['user_id'].unique()
        return len(long_interval_users)

    def get_user_purchase_history(self, user_id):
        user_data = self.data[
            (self.data['user_id'] == user_id) &
            (self.data['is_buy'] == True)
            ].sort_values('datetime')

        if user_data.empty:
            return pd.DataFrame(columns=['购买时间', '上次购买时间', '时间间隔(天)'])

        user_data['上次购买时间'] = user_data['datetime'].shift(1)
        user_data['时间间隔(天)'] = (user_data['datetime'] - user_data['上次购买时间']).dt.days

        return user_data[['datetime', '上次购买时间', '时间间隔(天)']].rename(
            columns={'datetime': '购买时间'}
        )

    def get_behavior_distribution(self):
        behavior_counts = self.data['behavior_type'].value_counts().reset_index()
        behavior_counts.columns = ['行为类型', '次数']
        behavior_counts['比例'] = (behavior_counts['次数'] / behavior_counts['次数'].sum()).map(lambda x: f"{x:.2%}")
        return behavior_counts

    def get_user_active_hours(self):
        hourly_activity = self.data.groupby('hour')['user_id'].count().reset_index(name='行为次数')
        return hourly_activity

    def get_user_active_days(self):
        day_names = ['周一', '周二', '周三', '周四', '周五', '周六', '周日']
        daily_activity = self.data.groupby('day_of_week')['user_id'].count().reset_index()
        daily_activity.columns = ['day_of_week', '行为次数']
        daily_activity['周别'] = daily_activity['day_of_week'].map(lambda x: day_names[x])
        return daily_activity[['周别', '行为次数']]

    def get_top_viewed_items(self, n=10):
        viewed_items = self.data[self.data['behavior_type'] == 'pv'].groupby('item_id')['user_id'].count().reset_index(
            name='浏览量')
        return viewed_items.sort_values('浏览量', ascending=False).head(n)

    def get_top_purchased_items(self, n=10):
        purchased_items = self.data[self.data['behavior_type'] == 'buy'].groupby('item_id')[
            'user_id'].count().reset_index(name='购买量')
        return purchased_items.sort_values('购买量', ascending=False).head(n)

    def get_view_and_purchase_correlation(self):
        item_stats = self.data.groupby('item_id')['behavior_type'].value_counts().unstack(fill_value=0)
        item_stats = item_stats.rename(columns={'pv': '浏览量', 'buy': '购买量'})
        correlation = item_stats[['浏览量', '购买量']].corr().iloc[0, 1]

        return {
            'item_stats': item_stats,
            'correlation': correlation
        }